package com.txc.common.utils.jwt;

import cn.hutool.core.codec.Base64Decoder;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.RSA;
import lombok.extern.slf4j.Slf4j;

import java.io.*;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA加密解密工具类
 *
 * @author tao
 */
@Slf4j
public class RSAUtils {
    private static RSA rsa;
    private static final String PUBLIC_KEY_FILE_NAME = "static/pem/public.pem";
    private static final String PRIVATE_KEY_FILE_NAME = "static/pem/private.pem";
    private static final String PUBLIC_KEY_TYPE = "PUBLIC KEY";
    private static final String PRIVATE_KEY_TYPE = "PRIVATE KEY";
    private static final String RSA_ALGORITHM = "RSA";

    static {
        try {
            rsa = new RSA();
            rsa.setPublicKey(getPublicKeyFromPem());
            rsa.setPrivateKey(getPrivateKeyFromPem());
        } catch (Exception e) {
            log.error("RSAUtils init error", e);
        }
    }

    /**
     * 加密
     *
     * @param data 需要加密数据
     * @return 加密后数据 String
     */
    public static String encrypt(String data) {
        return rsa.encryptBase64(data, KeyType.PublicKey);
    }

    /**
     * 解密
     *
     * @param data 需要解密数据
     * @return 解密后数据 String
     */
    public static String decrypt(String data) {
        return rsa.decryptStr(data, KeyType.PrivateKey);
    }

    /**
     * 加密
     *
     * @param data 需要加密数据
     * @return 加密后数据 byte[]
     */
    public static byte[] encrypt(byte[] data) {
        return rsa.encrypt(data, KeyType.PublicKey);
    }

    /**
     * 解密
     *
     * @param data 需要解密数据
     * @return 解密后数据 byte[]
     */
    public static byte[] decrypt(byte[] data) {
        return rsa.decrypt(data, KeyType.PrivateKey);
    }

    /**
     * 读取私钥文件
     *
     * @return 私钥
     */
    private static PrivateKey getPrivateKeyFromPem() {
        try {
            String str = getResourceFileContent(PRIVATE_KEY_FILE_NAME);
            byte[] b = Base64Decoder.decode(str);
            KeyFactory kf = KeyFactory.getInstance(RSA_ALGORITHM);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(b);
            return kf.generatePrivate(keySpec);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 读取公钥文件
     *
     * @return 公钥
     */
    private static PublicKey getPublicKeyFromPem() {
        try {
            String str = getResourceFileContent(PUBLIC_KEY_FILE_NAME);
            byte[] b = Base64Decoder.decode(str);
            KeyFactory kf = KeyFactory.getInstance(RSA_ALGORITHM);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(b);
            return kf.generatePublic(keySpec);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 读取资源文件
     *
     * @param fileName 文件名
     * @return 文件内容
     */
    private static String getResourceFileContent(String fileName) {
        InputStream publicKeyStream = RSAUtils.class.getClassLoader().getResourceAsStream(fileName);
        if (publicKeyStream == null) {
            throw new RuntimeException("file not found: " + fileName);
        }
        try (BufferedReader br = new BufferedReader(new InputStreamReader(publicKeyStream))) {
            String line;
            StringBuilder str = new StringBuilder();
            while (ObjectUtil.isNotEmpty((line = br.readLine()))) {
                if (line.startsWith("-")) {
                    continue;
                }
                str.append(line);
            }
            return str.toString();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }


    public static void generateKeyPair(String publicKeyFileName, String privateKeyFileName, Integer keySize) throws IOException {
        KeyPair keyPair = SecureUtil.generateKeyPair(RSA_ALGORITHM, keySize);
        savePemFile(keyPair.getPrivate(), privateKeyFileName);
        savePemFile(keyPair.getPublic(), publicKeyFileName);

    }

    private static void savePemFile(Key key, String filename)
            throws IOException {
        String encodedKey = Base64.getEncoder().encodeToString(key.getEncoded());
        String keyType = (key instanceof PrivateKey) ? PRIVATE_KEY_TYPE : PUBLIC_KEY_TYPE;
        try (FileWriter writer = new FileWriter(filename)) {
            writer.write("-----BEGIN " + keyType + "-----\n");
            writer.write(encodedKey);
            writer.write("\n-----END " + keyType + "-----");
            writer.flush();
        }
    }

}
